{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "30836ad1",
   "metadata": {},
   "source": [
    "# 简单实现 CEPA 方法"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076eda91",
   "metadata": {},
   "source": [
    "> 创建时间：2021-09-04"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "867b57cd",
   "metadata": {},
   "source": [
    "本文档会简单地实现闭壳层 CEPA($n$) 方法。\n",
    "\n",
    "CEPA($n$) 不太完整地说，是基于 CCSD 与 CISD 的近似。一般认为这类方法是严格的 CCSD 近似，对 MP2 甚至 MP3 有所提升；但依据推导方式不同，可以是 CISD 的近似，也可以是 CISD 的补充。CEPA($n$) 方法的相对完善的综述可以参考 [^Ahlrichs.CPC.1979]。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "71b95aa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pyscf import gto, scf, ci, lib\n",
    "from pyscf.cc import ccsd\n",
    "from opt_einsum import contract as einsum\n",
    "\n",
    "np.set_printoptions(6, suppress=True, linewidth=120)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bb51087",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fcbe9a0e",
   "metadata": {},
   "source": [
    "## CISD 能量"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f0c6cdd",
   "metadata": {},
   "source": [
    "### 准备工作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22656d9b",
   "metadata": {},
   "source": [
    "我们以键长 0.96 Angstrom、键角 104.5° 的水分子作为研究对象。基组是 cc-pVDZ。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f34e0f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = gto.Mole(atom=\"O; H 1 0.96; H 1 0.96 2 104.5\", basis=\"cc-pVDZ\", verbose=0).build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "50763ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "mf_scf = scf.RHF(mol).run()\n",
    "mf_ci = ci.CISD(mf_scf).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bcfa3b2",
   "metadata": {},
   "source": [
    "我们首先考察 CISD 能量的计算。回顾到 CISD 计算方程组是\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "(\\hat H - E_\\textsf{HF} - E_\\mathrm{c}^\\textsf{CISD}) | \\Psi \\rangle = 0\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c396b503",
   "metadata": {},
   "source": [
    "其中，我们规定 $|\\Psi\\rangle$ 是半归一化 (Intermediate Normalized) 的波函数：\n",
    "\n",
    "$$\n",
    "|\\Psi\\rangle = |\\Phi_0\\rangle + \\sum_{ia} t_i^a |\\Phi_i^a\\rangle + \\sum_{ijab} t_{ij}^{ab} |\\Phi_{ij}^{ab}\\rangle\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53874596",
   "metadata": {},
   "source": [
    "$|\\Phi_0\\rangle$ 是 HF 基态波函数，$|\\Phi_i^a\\rangle$ 是单激发波函数、$|\\Phi_{ij}^{ab}\\rangle$ 是双激发波函数。需要强调的是，由于是空间轨道，因此这些波函数并不反映物理实在，特别是我们无法考察自旋算符 $\\hat S{}^2$ 的本征态。它们只是用来计算能量的方便的记号。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd82a9f",
   "metadata": {},
   "source": [
    "若要针对一次激发、二次激发分别给出 CISD 方程，那么\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\langle \\Phi_i^a | \\hat H - E_\\textsf{HF} - E_\\mathrm{c}^\\textsf{CISD} | \\Psi \\rangle &= 0 \\\\\n",
    "\\langle \\Phi_{ij}^{ab} | \\hat H - E_\\textsf{HF} - E_\\mathrm{c}^\\textsf{CISD} | \\Psi \\rangle &= 0 \\\\\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bb943ad",
   "metadata": {},
   "source": [
    "### 激发系数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59601de6",
   "metadata": {},
   "source": [
    "通过 PySCF 的 `ci` attribute，进而通过 `cisdvec_to_amplitudes` 函数可以得到 $|\\Psi\\rangle$ 的激发系数，但该激发系数并不是归一化的。我们手动将其进行半归一化，得到系数 `t1` $t_i^a$ 与 `t2` $t_{ij}^{ab}$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ec9e9702",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((5, 19), (5, 5, 19, 19))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "civec = mf_ci.ci / mf_ci.ci[0]\n",
    "_, t1, t2 = mf_ci.cisdvec_to_amplitudes(civec)\n",
    "t1.shape, t2.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9aca702",
   "metadata": {},
   "source": [
    "### 能量表达式验证"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "746a89a6",
   "metadata": {},
   "source": [
    "首先我们要给出电子积分。在有足够内存的情况下，可以用 `ccsd` 的 `_make_eris_incore` 实现："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3b85fded",
   "metadata": {},
   "outputs": [],
   "source": [
    "eris = ccsd._make_eris_incore(mf_ci, mf_scf.mo_coeff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f141701",
   "metadata": {},
   "source": [
    "我们通篇文档不关心涉及 $\\hat H$ 的计算过程。我们暂时只需要其中的\n",
    "\n",
    "$$\n",
    "g_{ij}^{ab} = \\langle ij | ab \\rangle = (ia|jb) = \\iint \\phi_i(\\boldsymbol{r}_1) \\phi_a(\\boldsymbol{r}_1) \\frac{1}{|\\boldsymbol{r}_1 - \\boldsymbol{r}_2|} \\phi_j(\\boldsymbol{r}_2) \\phi_b(\\boldsymbol{r}_2) \\, \\mathrm{d} \\boldsymbol{r}_1 \\, \\mathrm{d} \\boldsymbol{r}_2\n",
    "$$\n",
    "\n",
    "该张量可以通过 `eris.ovov` 调出，维度是 $(i, a, j, b)$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "549b3ecc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5, 19, 5, 19)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eris.ovov.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02ced1bf",
   "metadata": {},
   "source": [
    "对于半归一化方法，不论是 MP2, CEPA($n$), CI(S)D 或 CC(S)D，下式在闭壳层下总是成立的：\n",
    "\n",
    "$$\n",
    "E_\\mathrm{c} = \\sum_{ijab} (2 t_{ij}^{ab} - t_{ij}^{ba}) g_{ij}^{ab}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b284cb07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.2053384394211093"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2 * einsum(\"iajb, ijab ->\", eris.ovov, t2) - einsum(\"iajb, ijba ->\", eris.ovov, t2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf548d82",
   "metadata": {},
   "source": [
    "我们再次回顾 PySCF 计算得到的 CISD 相关能是下述非常接近的结果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f6279172",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.20533844297533488"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_ci.e_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f393ccb2",
   "metadata": {},
   "source": [
    "## CEPA($n$) 计算"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18b958b",
   "metadata": {},
   "source": [
    "### CEPA($n$) 原理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30af7e5d",
   "metadata": {},
   "source": [
    "CEPA($n$) 方程组是\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\langle \\Phi_{i}^{a} | \\hat H - E_\\textsf{HF} - B_{i} | \\Psi \\rangle &= 0 \\\\\n",
    "\\langle \\Phi_{ij}^{ab} | \\hat H - E_\\textsf{HF} - A_{ij} | \\Psi \\rangle &= 0\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c76f0c61",
   "metadata": {},
   "source": [
    "其中，\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "A_{ij} = \\left\\{\n",
    "\\begin{matrix}\n",
    "0 & \\textsf{CEPA(0)} \\\\\n",
    "e_{ij} & \\textsf{CEPA(2)} \\\\\n",
    "\\frac{1}{2} \\sum_k (e_{ik} + e_{kj}) & \\textsf{CEPA(1)} \\\\\n",
    "\\sum_k (e_{ik} + e_{kj}) - e_{ij} & \\textsf{CEPA(3)} \\\\\n",
    "E_\\mathrm{c}^\\textsf{CISD} = \\sum_{kl} e_{kl} & \\textsf{CISD}\n",
    "\\end{matrix}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52539f63",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "B_{i} = \\left\\{\n",
    "\\begin{matrix}\n",
    "0 & \\textsf{CEPA(0)} \\\\\n",
    "\\textsf{NaN} & \\textsf{CEPA(2)} \\\\\n",
    "\\sum_k e_{ik} & \\textsf{CEPA(1)} \\\\\n",
    "2 \\sum_k e_{ik} - e_{ii} & \\textsf{CEPA(3)} \\\\\n",
    "E_\\mathrm{c}^\\textsf{CISD} = \\sum_{kl} e_{kl} & \\textsf{CISD}\n",
    "\\end{matrix}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a9741c",
   "metadata": {},
   "source": [
    "对电子能定义为\n",
    "\n",
    "$$\n",
    "e_{ij} = \\sum_{ab} (2 t_{ij}^{ab} - t_{ij}^{ba}) g_{ij}^{ab}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1c49c257",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pair_energy(eris_ovov, t2):\n",
    "    return 2 * einsum(\"iajb, ijab -> ij\", eris_ovov, t2) - einsum(\"iajb, ijba -> ij\", eris_ovov, t2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eb2b14e",
   "metadata": {},
   "source": [
    "CEPA(2) 由于不确定其 $B_i$ 的定义，因此这里不作实现。当然，CEPA(3) 的定义也可能存在疑问。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "84a52f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cepa_shift(cepa_n, t1, t2, e_ij):\n",
    "    e_i, e_j = e_ij.sum(axis=1), e_ij.sum(axis=0)\n",
    "    A_ij, B_i = np.zeros_like(e_ij), np.zeros_like(e_i)\n",
    "    if cepa_n == 0:\n",
    "        pass\n",
    "    elif cepa_n == 1:\n",
    "        A_ij = 0.5 * (e_i[:, None] + e_j[None, :])\n",
    "        B_i = e_i\n",
    "    elif cepa_n == 3:\n",
    "        A_ij = e_i[:, None] + e_j[None, :] - e_ij\n",
    "        B_i = - e_ij.diagonal() + 2 * e_i\n",
    "    else:\n",
    "        raise ValueError(\"cepa_n value error for \" + cepa_n)\n",
    "    return A_ij, B_i"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2138dae9",
   "metadata": {},
   "source": [
    "### 迭代法求取 CISD 系数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f41b895",
   "metadata": {},
   "source": [
    "我们首先考虑二次激发系数 $t_{ij}^{ab}$ 的求取。\n",
    "\n",
    "$$\n",
    "\\langle \\Phi_{ij}^{ab} | \\hat H - E_\\textsf{HF} | \\Psi \\rangle = A_{ij} \\langle \\Phi_{ij}^{ab} | \\Psi \\rangle = A_{ij} t_{ij}^{ab}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "137b2137",
   "metadata": {},
   "source": [
    "如果我们记等式左为 $\\mathscr{v}_{ij}^{ab}$，那么该式整理为\n",
    "\n",
    "$$\n",
    "\\mathscr{v}_{ij}^{ab} = A_{ij} t_{ij}^{ab} \\quad \\Rightarrow \\quad t_{ij}^{ab} = \\frac{\\mathscr{v}_{ij}^{ab}}{A_{ij}}\n",
    "$$\n",
    "\n",
    "但需要注意，这不是一个好的激发系数 $t_{ij}^{ab}$ 的更新策略。这里指出，由于 $\\langle \\Phi_{ij}^{ab} | \\hat H - E_\\textsf{HF} | \\Phi_0 \\rangle = -D_{ij}^{ab}$ 是 $\\mathscr{v}_{ij}^{ab}$ 的重要贡献项，因此不妨将上式写为\n",
    "\n",
    "$$\n",
    "\\mathscr{v}_{ij}^{ab} - D_{ij}^{ab} t_{ij}^{ab} + D_{ij}^{ab} t_{ij}^{ab} = A_{ij} t_{ij}^{ab} \\quad \\Rightarrow \\quad t_{ij}^{ab} = t_{ij}^{ab} + \\frac{\\mathscr{v}_{ij}^{ab} - A_{ij} t_{ij}^{ab}}{D_{ij}^{ab}}\n",
    "$$\n",
    "\n",
    "上式是我们实际会使用到的激发系数 $t_{ij}^{ab}$ 的更新策略。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf138b25",
   "metadata": {},
   "source": [
    "在迭代计算激发系数时，我们需要给定初始激发系数。一般来说使用 MP2 激发系数即可：\n",
    "\n",
    "$$\n",
    "\\tilde t_i^a = 0, \\quad \\tilde t_{ij}^{ab} = g_{ij}^{ab} / D_{ij}^{ab}\n",
    "$$\n",
    "\n",
    "其中，`d1` $D_i^a = \\varepsilon_i - \\varepsilon_a$，`d2` $D_{ij}^{ab} = \\varepsilon_i + \\varepsilon_j - \\varepsilon_a - \\varepsilon_b$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f0803ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def corr_cepa(cepa_n):\n",
    "    # Prepare D_i^a, D_ij^ab\n",
    "    d0, d1, d2 = mf_ci.cisdvec_to_amplitudes(mf_ci.make_diagonal(eris))\n",
    "    d1 = d0 - d1; d2 = d0 - d2\n",
    "    # Prepare initial t_i^a, t_ij^ab\n",
    "    t0, t1, t2 = 1, np.zeros_like(d1), eris.ovov.swapaxes(1, 2) / d2\n",
    "    civec = mf_ci.amplitudes_to_cisdvec(t0, t1, t2)\n",
    "    # Iteration\n",
    "    for it in range(100):\n",
    "        t1_old, t2_old = t1, t2\n",
    "        _, v1, v2 = mf_ci.cisdvec_to_amplitudes(mf_ci.contract(civec, eris))\n",
    "        e_ij = pair_energy(eris.ovov, t2)\n",
    "        A_ij, B_i = cepa_shift(cepa_n, t1, t2, e_ij)\n",
    "        t1 = t1_old + (v1 - B_i[:, None]           * t1_old) / d1\n",
    "        t2 = t2_old + (v2 - A_ij[:, :, None, None] * t2_old) / d2\n",
    "        e_corr = pair_energy(eris.ovov, t2).sum()\n",
    "        civec = mf_ci.amplitudes_to_cisdvec(t0, t1, t2)\n",
    "        if np.linalg.norm(t1 - t1_old) + np.linalg.norm(t2 - t2_old) < 1e-5:\n",
    "            print(\"Total Iterations: \", it)\n",
    "            print(\"Correlation Energy: \", e_corr)\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6ffd2004",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Iterations:  26\n",
      "Correlation Energy:  -0.2167752177602909\n"
     ]
    }
   ],
   "source": [
    "corr_cepa(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "86d942a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Iterations:  38\n",
      "Correlation Energy:  -0.21352333911480398\n"
     ]
    }
   ],
   "source": [
    "corr_cepa(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eaa8779f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Iterations:  53\n",
      "Correlation Energy:  -0.21129784897010107\n"
     ]
    }
   ],
   "source": [
    "corr_cepa(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab8a9c0",
   "metadata": {},
   "source": [
    "## 参考文献"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "678b5b2b",
   "metadata": {},
   "source": [
    "[^Ahlrichs.CPC.1979]: Ahlrichs, R. Many Body Perturbation Calculations and Coupled Electron Pair Models. *Comput. Phys. Commun.* **1979**, *17* (1), 31–45. doi: [10.1016/0010-4655(79)90067-5](https://doi.org/10.1016/0010-4655(79)90067-5)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
